package com.hk.core.data.r2dbc.tenants;

import com.hk.commons.util.Lazy;
import com.hk.commons.util.SpringContextHolder;
import com.hk.commons.util.StringUtils;
import com.hk.core.data.r2dbc.annotations.TenantId;
import io.r2dbc.spi.Row;
import io.r2dbc.spi.RowMetadata;
import lombok.RequiredArgsConstructor;
import org.springframework.dao.DataAccessException;
import org.springframework.data.r2dbc.convert.R2dbcConverter;
import org.springframework.data.r2dbc.core.R2dbcEntityOperations;
import org.springframework.data.r2dbc.core.ReactiveDataAccessStrategy;
import org.springframework.data.relational.core.query.Criteria;
import org.springframework.data.relational.core.query.Query;
import org.springframework.data.relational.core.query.Update;
import org.springframework.r2dbc.core.DatabaseClient;
import org.springframework.r2dbc.core.PreparedOperation;
import org.springframework.r2dbc.core.RowsFetchSpec;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Objects;
import java.util.function.BiFunction;

/**
 * @author Kevin
 * @date 2024-04-20 22:15
 */
@RequiredArgsConstructor
public class TenantR2dbcEntityTemplate implements R2dbcEntityOperations {

    private final R2dbcEntityOperations delegate;

    private final Lazy<UserTenantAware> userTenantAwareLazy = Lazy.of(() -> SpringContextHolder.getBean(UserTenantAware.class));

    @Override
    public DatabaseClient getDatabaseClient() {
        return delegate.getDatabaseClient();
    }

    @Override
    public ReactiveDataAccessStrategy getDataAccessStrategy() {
        return delegate.getDataAccessStrategy();
    }

    @Override
    public R2dbcConverter getConverter() {
        return delegate.getConverter();
    }

    /**
     * 实体必须要有 {@link TenantId } 注解标注的属性
     */
    private String getTenantColumn(Class<?> entityClass) {
        var persistentProperty = delegate.getConverter().getMappingContext()
                .getRequiredPersistentEntity(entityClass).getPersistentProperty(TenantId.class);
        return Objects.isNull(persistentProperty) ? null : persistentProperty.getColumnName().getReference();
    }

    private Query getTenantIdQuery(Query query, String tenantId, Class<?> entityClass) {
        String tenantColumn = getTenantColumn(entityClass);
        if (!StringUtils.hasLength(tenantColumn)) {
            return query;
        }
        return query.getCriteria().map(definition -> {
            String reference = Objects.nonNull(definition.getColumn()) ? definition.getColumn().getReference() : null;
            do {
                if (StringUtils.equals(reference, tenantColumn)) {//说明条件中已添加过tenant_id的条件，直接返回即可.
                    return query;
                }
                var previous = definition.getPrevious();
                reference = Objects.nonNull(previous) && Objects.nonNull(previous.getColumn()) ? previous.getColumn().getReference() : null;
            } while (StringUtils.hasLength(reference));
            return Query.query(Criteria.where(tenantColumn).is(tenantId).and(definition));// 添加一个tenant_id 的查询条件
        }).orElse(Query.query(Criteria.where(tenantColumn).is(tenantId))); //如果没有查询条件，也默认给一个tenant_id的查询条件
    }

    @Override
    public Mono<Long> count(Query query, Class<?> entityClass) throws DataAccessException {
        return userTenantAwareLazy.get().getCurrentTenantId()
                .flatMap(tenantId -> delegate.count(getTenantIdQuery(query, tenantId, entityClass), entityClass));
    }

    @Override
    public Mono<Boolean> exists(Query query, Class<?> entityClass) throws DataAccessException {
        return userTenantAwareLazy.get().getCurrentTenantId()
                .flatMap(tenantId -> delegate.exists(getTenantIdQuery(query, tenantId, entityClass), entityClass));
    }

    @Override
    public <T> Flux<T> select(Query query, Class<T> entityClass) throws DataAccessException {
        return userTenantAwareLazy.get().getCurrentTenantId()
                .flatMapMany(tenantId -> delegate.select(getTenantIdQuery(query, tenantId, entityClass), entityClass));
    }

    @Override
    public <T> Mono<T> selectOne(Query query, Class<T> entityClass) throws DataAccessException {
        return userTenantAwareLazy.get().getCurrentTenantId()
                .flatMap(tenantId -> delegate.selectOne(getTenantIdQuery(query, tenantId, entityClass), entityClass));
    }

    @Override
    public Mono<Long> update(Query query, Update update, Class<?> entityClass) throws DataAccessException {
        return delegate.update(query, update, entityClass);
    }

    @Override
    public Mono<Long> delete(Query query, Class<?> entityClass) throws DataAccessException {
        return delegate.delete(query, entityClass);
    }

    @Override
    public <T> RowsFetchSpec<T> query(PreparedOperation<?> operation, Class<T> entityClass) throws DataAccessException {
        return delegate.query(operation, entityClass);
    }

    @Override
    public <T> RowsFetchSpec<T> query(PreparedOperation<?> operation, Class<?> entityClass, Class<T> resultType) throws DataAccessException {
        return delegate.query(operation, entityClass, resultType);
    }

    @Override
    public <T> RowsFetchSpec<T> query(PreparedOperation<?> operation, BiFunction<Row, RowMetadata, T> rowMapper) throws DataAccessException {
        return delegate.query(operation, rowMapper);
    }

    @Override
    public <T> RowsFetchSpec<T> query(PreparedOperation<?> operation, Class<?> entityClass, BiFunction<Row, RowMetadata, T> rowMapper) throws DataAccessException {
        return delegate.query(operation, entityClass, rowMapper);
    }

    @Override
    public <T> RowsFetchSpec<T> getRowsFetchSpec(DatabaseClient.GenericExecuteSpec executeSpec, Class<?> entityType, Class<T> resultType) {
        return delegate.getRowsFetchSpec(executeSpec, entityType, resultType);
    }

    @Override
    public <T> Mono<T> insert(T entity) throws DataAccessException {
        return delegate.insert(entity);
    }

    @Override
    public <T> Mono<T> update(T entity) throws DataAccessException {
        return delegate.update(entity);
    }

    @Override
    public <T> Mono<T> delete(T entity) throws DataAccessException {
        return delegate.delete(entity);
    }

    @Override
    public ReactiveDelete delete(Class<?> domainType) {
        return delegate.delete(domainType);
    }

    @Override
    public <T> ReactiveInsert<T> insert(Class<T> domainType) {
        return delegate.insert(domainType);
    }

    @Override
    public <T> ReactiveSelect<T> select(Class<T> domainType) {
        return delegate.select(domainType);
    }

    @Override
    public ReactiveUpdate update(Class<?> domainType) {
        return delegate.update(domainType);
    }
}
